/*
* Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
* Description: Hccl One-Sided Service
* Author: yanzhenghang
* Create: 2024-09-06
*/

#ifndef HCCL_ONE_SIDED_SERVICE_H
#define HCCL_ONE_SIDED_SERVICE_H

#include "i_hccl_one_sided_service.h"
#include "hccl_one_sided_conn.h"
#include "rma_buffer_mgr.h"
#include "local_rdma_rma_buffer.h"
#include "local_ipc_rma_buffer.h"
#include "hccl_common.h"
#include "externalinput.h"

using HcclBatchData = struct HcclBatchDataDef {
    HcclComm comm;
    HcclCMDType cmdType;
    u32 remoteRank;
    HcclOneSideOpDesc* desc;
    u32 descNum;
    rtStream_t stream;
};

namespace hccl {
class HcclOneSidedService : public IHcclOneSidedService {
public:
    using RankId = u32;
    using LocalRdmaRmaBufferMgr = typename HcclOneSidedConn::LocalRdmaRmaBufferMgr;
    using LocalIpcRmaBufferMgr = typename HcclOneSidedConn::LocalIpcRmaBufferMgr;
    using ProcessInfo = typename HcclOneSidedConn::ProcessInfo;
    using MemDescData = typename HcclOneSidedConn::MemDescData;

    struct MemInfo {
        void *addr{nullptr};
        u64 size{0};
        HcclMemType type{HcclMemType::HCCL_MEM_TYPE_DEVICE};
        MemInfo(void *addr, u64 size, HcclMemType type)
          : addr(addr), size(size), type(type) {}
    };

    HcclOneSidedService(std::unique_ptr<HcclSocketManager> &socketManager,
        std::unique_ptr<NotifyPool> &notifyPool);

    // 父类Config()等已经完成必要参数的配置
    HcclOneSidedService() = default;
    ~HcclOneSidedService() = default;

    HcclResult RegMem(void* addr, u64 size, HcclMemType type, RankId remoteRankId, HcclMemDesc &localMemDesc);
    HcclResult DeregMem(const HcclMemDesc &localMemDesc);

    // 可能返回超时
    HcclResult ExchangeMemDesc(RankId remoteRankId, const HcclMemDescs &localMemDescs,
        HcclMemDescs &remoteMemDescs, u32 &actualNumOfRemote, const std::string &commIdentifier, s32 timeoutSec);

    void EnableMemAccess(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem);
    void DisableMemAccess(const HcclMemDesc &remoteMemDesc);

    void BatchPut(RankId remoteRankId, const HcclOneSideOpDesc* desc, u32 descNum, const rtStream_t &stream);
    void BatchGet(RankId remoteRankId, const HcclOneSideOpDesc* desc, u32 descNum, const rtStream_t &stream);

private:
    u32 registedMemCnt_{0};
    HcclResult IsUsedRdma(RankId remoteRankId, bool &useRdma);
    void RegMemRdma(MemInfo &memInfo, RankId remoteRankId, HcclMemDesc &localMemDesc, bool &isRegistedSucc);
    void RegMemIpc(MemInfo &memInfo, RankId remoteRankId, HcclMemDesc &localMemDesc, bool &isRegistedSucc);
    void DeregMemRdma(const HcclMemDesc &localMemDesc, bool &isDeregistedSucc);
    void DeregMemIpc(const HcclMemDesc &localMemDesc, bool &isDeregistedSucc);
    HcclResult SetupRemoteRankInfo(RankId remoteRankId, HcclRankLinkInfo &remoteRankInfo);
    HcclResult CreateConnection(RankId remoteRankId, const HcclRankLinkInfo &remoteRankInfo,
        std::shared_ptr<HcclOneSidedConn> &tempConn, const std::string &commIdentifier, s32 timeoutSec);
    HcclResult Grant(const HcclMemDesc &localMemDesc, const ProcessInfo &remoteProcess);

    std::unordered_map<RankId, std::shared_ptr<HcclOneSidedConn>> oneSidedConns_{};
    std::unordered_map<std::string, std::shared_ptr<LocalRdmaRmaBuffer>> desc2LocalRdmaRmaBufferMap_{};
    std::unordered_map<std::string, std::shared_ptr<LocalIpcRmaBuffer>> desc2LocalIpcRmaBufferMap_{};

    std::unordered_map<RankId, bool> isUsedRdmaMap_;
    LocalRdmaRmaBufferMgr localRdmaRmaBufferMgr_{};
    LocalIpcRmaBufferMgr localIpcRmaBufferMgr_{};
};
}

#endif